import cv2
import numpy as np
import os
import matplotlib.pyplot as plt


def process_image(input_image_path, output_folder='chars'):
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 读取图像
    img = cv2.imread(input_image_path, cv2.IMREAD_GRAYSCALE)

    # 全局阈值处理
    _, binary_img = cv2.threshold(img, 128, 255, cv2.THRESH_BINARY_INV)

    # 腐蚀操作
    kernel1 = np.ones((3, 3), np.uint8)
    eroded_img = cv2.erode(binary_img, kernel1, iterations=4)

    # 膨胀操作
    kernel2 = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    dilated_img = cv2.dilate(eroded_img, kernel2, iterations=5)
    dilated_img = cv2.medianBlur(dilated_img, 5)

    # 应用闭合运算
    kernel3 = np.ones((5, 5), np.uint8)
    closed_img = cv2.morphologyEx(dilated_img, cv2.MORPH_CLOSE, kernel3, iterations=10)

    # Canny边缘检测
    edges = cv2.Canny(closed_img, 100, 200)

    # 获取轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    # 将灰度图像转换为彩色图像，以绘制绿色框
    img_contours = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    for c in contours:
        perimeter = cv2.arcLength(c, True)
        if perimeter > 100:
            x, y, w, h = cv2.boundingRect(c)
            cv2.rectangle(img_contours, (x, y), (x + w, y + h), (0, 255, 0), 2)

    # 保存处理结果
    cv2.imwrite(os.path.join(output_folder, 'original_gray.jpg'), img)
    cv2.imwrite(os.path.join(output_folder, 'binary_img.jpg'), binary_img)
    cv2.imwrite(os.path.join(output_folder, 'eroded_img.jpg'), eroded_img)
    cv2.imwrite(os.path.join(output_folder, 'dilated_img.jpg'), dilated_img)
    cv2.imwrite(os.path.join(output_folder, 'closed_img.jpg'), closed_img)
    cv2.imwrite(os.path.join(output_folder, 'edges.jpg'), edges)
    cv2.imwrite(os.path.join(output_folder, 'img_contours.jpg'), img_contours)

    # 显示处理结果
    images = [img, binary_img, eroded_img, dilated_img, closed_img, edges, img_contours]
    titles = ['原图', '二值图', '腐蚀', '膨胀', '闭合运算', '边缘检测', '结果']

    plt.figure(figsize=(15, 10))
    for i in range(len(images)):
        plt.subplot(2, 4, i + 1)
        if len(images[i].shape) == 2:  # 如果是灰度图像
            plt.imshow(images[i], cmap='gray')
        else:  # 如果是彩色图像
            plt.imshow(cv2.cvtColor(images[i], cv2.COLOR_BGR2RGB))
        plt.title(titles[i])
        plt.axis('off')

    plt.tight_layout()
    plt.savefig(os.path.join(output_folder, 'combined_result.png'))
    plt.show()


# 使用示例
input_image_path = 'hanzi1.jpg'
process_image(input_image_path)
